#include "BarDetector.h"

namespace hitcrt {
/**
 * @brief 接口默认实现
 * @param[in] roi           识别范围框
 * @param[in] mask          预处理过的二值图
 * @param[in] enemyChannel  敌方通道
 * @param[in] selfChannel   己方通道
 * @param[out] lightBars     输出灯条
 * @return true 检测到灯条
 * @return false
 */
bool BarDetectorBase::detect(const ROI roi, const cv::Mat& mask,
                             const cv::Mat& enemyChannel,
                             const cv::Mat& selfChannel,
                             std::vector<LightBar>& lightBars) {
    // 清空vectors成员变量
    clearVectors();
    // 找轮廓
    if (!findContours(mask)) {
        // std::cout << "无轮廓" << std::endl;
        return false;
    } else {
        // 找灯条
        //std::cout << "轮廓数量为：" << m_contours.size() << std::endl;
        return findLightBars(roi, enemyChannel, selfChannel, lightBars);
    }
}

/**
 * @brief 清空vector成员变量
 */
void BarDetectorBase::clearVectors() {
/// TODO 需要填写 ///
#pragma omp parallel for num_threads(12)
    m_contours.clear();
    m_contours.shrink_to_fit();
    m_insidePointsVec.clear();
    m_insidePointsVec.shrink_to_fit();
    m_boundingRects.clear();
    m_boundingRects.shrink_to_fit();
    m_rotatedRects.clear();
    m_rotatedRects.shrink_to_fit();
}

/**
 * @brief 找轮廓
 * @param[in] mask   二值图
 * @return true 轮廓数量符合要求
 * @return false
 * @note 如果环境较复杂，存在很多小连通域时，该项会变很大
 */
bool BarDetectorBase::findContours(const cv::Mat& mask) {
    std::vector<cv::Vec4i> hierarchy;
    cv::findContours(mask, m_contours, hierarchy, cv::RETR_LIST,
                     cv::CHAIN_APPROX_SIMPLE);
    return isInRange<int>(m_contours.size(), m_contourSizeMin,
                          m_contourSizeMax);
}

/**
 * @brief 找灯条
 * @param[in] roi           识别范围框
 * @param[in] enemyChannel  敌方通道
 * @param[in] selfChannel   己方通道
 * @param[out] lightBars     灯条
 * @return true 找到灯条
 * @return false
 */
bool BarDetectorBase::findLightBars(const ROI roi, const cv::Mat& enemyChannel,
                                    const cv::Mat& selfChannel,
                                    std::vector<LightBar>& lightBars) {
    const int sizeBeforeDetect = lightBars.size();
    // 遍历轮廓找灯条
    int i = 0;
    for (auto contourIter : m_contours) {
        i++;
        LightBar bar;
        //std::cout << "第" << i << "个轮廓检验结果" << std::endl;
        if (!isLightBar(contourIter, enemyChannel, selfChannel, bar)) {
            continue;
        }

        getPoints(contourIter, roi, bar);

        lightBars.emplace_back(std::move(bar));
        //std::cout<<std::endl;
    }
    const int sizeAfterDetect = lightBars.size();
    return (sizeAfterDetect - sizeBeforeDetect > 0);
}

/**
 * @brief 灯条特征判断该轮廓是否为灯条
 * @param[in] contour       轮廓
 * @param[in] enemyChannel  敌方通道
 * @param[in] selfChannel   己方通道
 * @param[out] bar           灯条
 * @return true
 * @return false
 */
bool BarDetectorBase::isLightBar(const std::vector<cv::Point>& contour,
                                 const cv::Mat& enemyChannel,
                                 const cv::Mat& selfChannel, LightBar& bar) {
    // 几何特征检验
    if (!checkGeometryFeature(contour, bar)) {
        //std::cout << "由于几何特征 检验失败" << std::endl;
        return false;
    }

    // 颜色检验
    if (!checkColor(contour, enemyChannel, selfChannel, bar)) {
        //std::cout << "由于颜色检验 检验失败" << std::endl;

        return false;
    }
    // 方向检验
    if (!checkDirection(bar)) {
        //std::cout << "由于方向检验 检验失败" << std::endl;

        return false;
    }
    //std::cout << "检验通过" << std::endl;
    return true;
}

/**
 * @brief 灯条几何特征检验
 * @param[in] contour       轮廓
 * @param[out] bar           灯条
 * @return true 符合几何特征
 * @return false
 */
bool BarDetector::checkGeometryFeature(const std::vector<cv::Point>& contour,
                                       LightBar& bar) {
    if (checkGeometryFeatureImpl(contour, bar)) {
        m_boundingRects.emplace_back(bar.m_boundingRect);
        m_rotatedRects.emplace_back(bar.m_rotatedRect);
        return true;
    } else {
        return false;
    }
}

/**
 * @brief 面积、长宽等几何特征检验
 * @param[in] contour       轮廓
 * @param[out] bar           灯条
 * @return true
 * @return false
 */
bool BarDetector::checkGeometryFeatureImpl(
    const std::vector<cv::Point>& contour, LightBar& bar) {
    /// TODO 需要填写 ///
    // 灯条面积（pix）
    double area = cv::contourArea(contour);
    if (area < m_areaMin || area > m_areaMax) {
        //std::cout << "灯条面积不符" << std::endl;
        //std::cout << "灯条面积为：" << area << " " << "m_areaMin:" << m_areaMin
                  //<< " " << "m_areaMax:" << m_areaMax << std::endl;
        return false;
    }
    // 检验长、宽（pix）
    // 提取矩形
    cv::Rect Rect = cv::boundingRect(contour);
    cv::RotatedRect rotatedRect = cv::minAreaRect(contour);
    // 比较大小确定长和宽
    bar.m_length = rotatedRect.size.height > rotatedRect.size.width
                       ? rotatedRect.size.height
                       : rotatedRect.size.width;
    bar.m_width = rotatedRect.size.width > rotatedRect.size.height
                      ? rotatedRect.size.height
                      : rotatedRect.size.width;
    // 判断阈值关系
    if (bar.m_length < m_lengthMin || bar.m_length > m_lengthMax) {
        //std::cout << "灯条长度不符" << std::endl;
        //std::cout << "灯条长度为：" << bar.m_length << " "
                  //<< "m_lengthMin:" << m_lengthMin << " "
                  //<< "m_lengthMax:" << m_lengthMax << std::endl;
        return false;
    }
    // 图像上投影宽大于高要排除
    double aspectRatio = bar.m_width / bar.m_length;
    if (aspectRatio < m_aspectRatioMin || aspectRatio > m_aspectRatioMax) {
        //std::cout << "灯条宽长比不符" << std::endl;
        //std::cout << "灯条宽长比为：" << aspectRatio<< " "
                 // << "m_lengthMin:" << m_aspectRatioMin << " "
                 // << "m_lengthMax:" << m_aspectRatioMax << std::endl;
        return false;
    }
    bar.m_rotatedRect = rotatedRect;
    bar.m_boundingRect = Rect;

    bar.m_contourArea = area;
    return true;
}

/**
 * @brief 灯条颜色检验
 * @param[in] contour       轮廓
 * @param[in] enemyChannel  敌方通道
 * @param[in] selfChannel   己方通道
 * @param[out] bar           灯条
 * @return true 符合颜色特征
 * @return false
 */
bool BarDetector::checkColor(const std::vector<cv::Point>& contour,
                             const cv::Mat& enemyChannel,
                             const cv::Mat& selfChannel, LightBar& bar) {
    std::vector<cv::Point2f> insidePoints;
    if (checkColorImpl(contour, enemyChannel, selfChannel, insidePoints, bar)) {
        m_insidePointsVec.emplace_back(insidePoints);
        return true;
    } else {
        m_boundingRects.pop_back();
        m_rotatedRects.pop_back();
        return false;
    }
}

/**
 * @brief 颜色检验
 * @param  contour[i] 灯条轮廓
 * @param  enemyChannel[in] 敌方通道
 * @param  selfChannel[in]  我方图像通道
 * @param  insidePoints[out] 轮廓内满足条件的内点
 * @param  bar[out] 灯条
 * @return true
 * @return false
 * @note
 * 颜色检验，遍历轮廓内点判断其颜色是否符合要求
 */
bool BarDetector::checkColorImpl(const std::vector<cv::Point>& contour,
                                 const cv::Mat& enemyChannel,
                                 const cv::Mat& selfChannel,
                                 std::vector<cv::Point2f>& insidePoints,
                                 LightBar& bar) {
    /// TODO 需要填写
    /// 颜色检验，遍历轮廓判断像素点颜色差异是否符合要求，符合则放入内点///
    uchar colorDiffThresh, colorRatioThresh;
// 传入的图像应为单通道的值，例如对图片分离<R>通道和<B>通道
// 注意是  敌 - 我
#pragma omp parallel for num_threads(12)
    for (auto point : contour) {
        colorDiffThresh = enemyChannel.at<uchar>(point.y, point.x) -
                          selfChannel.at<uchar>(point.y, point.x);
        if (colorDiffThresh > m_colorDiffThresh) {
            insidePoints.push_back(cv::Point2f(point.x, point.y));
        }
    }
    //std::cout <<"内点数量为"<<insidePoints.size()<<std::endl;
    return insidePoints.size();
}

/**
 * @brief 方向检验
 * @param[in,out] bar 灯条
 * @return true 符合方向特征
 * @return false
 */
bool BarDetector::checkDirection(LightBar& bar) {
    if (checkDirectionImpl(m_insidePointsVec.back(), bar)) {
        return true;
    } else {
        m_insidePointsVec.pop_back();
        return false;
    }
}

/**
 * @brief 通过像素斜率检验灯条方向，必须在轮廓内点识别出来后使用
 * @param[in] insidePoints  轮廓内点
 * @param[out] bar           灯条
 * @return true 符合方向特征
 * @return false
 */
bool BarDetector::checkDirectionImpl(
    const std::vector<cv::Point2f>& insidePoints, LightBar& bar) {
    // 内点拟合直线,斜率限制,应该基本竖直向下
    /// TODO 需要填写 ///
    cv::Vec4f Line;
    cv::fitLine(insidePoints, Line, cv::DIST_L1, 0, 1e-2, 1e-2);
    m_line = Line;
    double k = std::fabs((double)Line[1] / (double)Line[0]);
    if (k < m_slopeThresh) {
        //std::cout<<"斜率检验不符"<<" "<<"K"<<k<<std::endl;
        //std::cout<<"m_slopeThresh: "<<m_slopeThresh<<std::endl;
        return false;
    }
    return true;
}

/**
 * @brief 获取灯条端点坐标
 * @param[in] contour       轮廓
 * @param[in] roi           识别范围框
 * @param[out] bar           灯条
 */
void BarDetector::getPoints(const std::vector<cv::Point>& contour,
                            const ROI roi, LightBar& bar) {
    getEndPointsImpl(contour, bar);
}

/**
 * @brief 检测轮廓端点
 * @param[in] contour       轮廓
 * @param[out] bar           灯条
 */
void BarDetector::getEndPointsImpl(const std::vector<cv::Point>& contour,
                                   LightBar& bar) {
    /// TODO 需要填写 ///

    if (m_line[0] == 0.0 && m_line[1] == 1.0) {
        cv::Point2f point1(bar.m_rotatedRect.center.x,
                           bar.m_rotatedRect.center.y + bar.m_length / 2);
        cv::Point2f point2(bar.m_rotatedRect.center.x,
                           bar.m_rotatedRect.center.y - bar.m_length / 2);
        bar.m_top = point1.y > point2.y ? point1 : point2;
        bar.m_bottom = point1.y < point2.y ? point1 : point2;
    } else {
        double dy = bar.m_length /
                    sqrt(m_line[0] * m_line[0] + m_line[1] * m_line[1]) *
                    m_line[1];
        double dx = bar.m_length /
                    sqrt(m_line[0] * m_line[0] + m_line[1] * m_line[1]) *
                    m_line[0];
        cv::Point2f point1(bar.m_rotatedRect.center.x - dx / 2,
                           bar.m_rotatedRect.center.y - dy / 2);
        cv::Point2f point2(bar.m_rotatedRect.center.x + dx / 2,
                           bar.m_rotatedRect.center.y + dy / 2);
        bar.m_top = point1.y > point2.y ? point1 : point2;
        bar.m_bottom = point1.y < point2.y ? point1 : point2;
    }
    bar.m_center = bar.m_rotatedRect.center;
    bar.m_directVec = bar.m_top - bar.m_bottom;
}

}  // namespace hitcrt
